/**
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 * CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

#ifndef COMM_MEM_MANAGER_H
#define COMM_MEM_MANAGER_H
#include <string>
#include <unordered_map>
#include <memory>
#include <mutex>
#include "hccl_api.h"
#include "mem_device_pub.h"
#include "ccl_buffer_manager.h"
#include "rma_buffer_mgr.h"

namespace hccl {

struct HcclMemHandle {
    void* addr {nullptr};
    uint64_t size {0};
    HcclMemType memType {HCCL_MEM_TYPE_DEVICE};
    HcclRegMemAttr attr {};
};

class CommMemMgr {
public:
    using Handle = std::shared_ptr<HcclMemHandle>;
    using MemKey = hccl::BufferKey<uintptr_t, uint64_t>;
    using Table  = hccl::RmaBufferMgr<MemKey, Handle>;
    CommMemMgr() = default;
    ~CommMemMgr() = default;

     // cclbuffer内存
    void CommSetHcclBufferSize(uint64_t size);
    HcclResult GetHcclBuffer(CommBuffer *buffer);
    
    // 用户注册/反注册内存
    HcclResult CommRegMem(const std::string& tag, const HcclMem& mem, HcclRegMemAttr attr, void** rawHandle);
    HcclResult CommUnregMem(const std::string& tag, const void* rawHandle);
    HcclResult CommGetLocalRegMemByTag(const std::string &tag, std::vector<HcclMem> &memVec);
private:
    static inline MemKey MakeKey(void* addr, uint64_t size) {
        return MemKey(reinterpret_cast<uintptr_t>(addr), static_cast<uint64_t>(size));
    }
    struct TagRegistry {
        Table table;                                        // 区间树 + ref 语义
    };

    // cclbuffer内存
    std::mutex bufferMutex_;
    uint32_t bufferSize_;
    CCLBufferManager bufferManager_;

    // 用户绑定内存
    std::mutex memMutex_;
    // 每个 tag 一份 registry
    std::unordered_map<std::string, TagRegistry> tagRegs_;
    // 每个tag n个 HcclMemHandle
    std::unordered_map<std::string, std::vector<std::shared_ptr<HcclMemHandle>>> opBindings_;
};
}


#endif